#!/usr/bin/env python3
"""
Firecrawl MCP 智能体 - 网页内容抓取
使用 Firecrawl MCP 服务来抓取和分析网页内容
"""

import asyncio
import os
from typing import Any, Dict

from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.tools import MCPTool
from langchain_community.llms import Tongyi
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import Tool
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 获取 Firecrawl API Key
FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY")

class FirecrawlMCPAgent:
    """Firecrawl MCP 智能体类"""
    
    def __init__(self):
        self.client = None
        self.tools = []
        self.agent = None
        self.llm = None
        
    async def setup_mcp_client(self):
        """设置 MCP 客户端连接"""
        print("正在设置 Firecrawl MCP 客户端...")
        print(f"使用 API Key: {FIRECRAWL_API_KEY[:10]}..." if FIRECRAWL_API_KEY else "未找到 API Key")
        
        # 配置 Firecrawl MCP 服务器连接（使用官方配置）
        server_config = {
            "transport": "stdio",
            "command": "npx",
            "args": ["-y", "firecrawl-mcp"],
            "env": {
                "FIRECRAWL_API_KEY": FIRECRAWL_API_KEY
            }
        }
        
        # 创建 MCP 客户端
        self.client = MultiServerMCPClient({
            "firecrawl-mcp": server_config
        })
        
        print("Firecrawl MCP 客户端设置完成")
        
    async def setup_tools(self):
        """设置工具"""
        print("正在设置工具...")
        
        try:
            # 建立 MCP 会话
            async with self.client.session("firecrawl-mcp") as session:
                # 获取可用工具
                tools_response = await session.list_tools()
                available_tools = [tool.name for tool in tools_response.tools]
                print(f"可用的 Firecrawl 工具: {available_tools}")
                
                # 创建 LangChain 工具
                for tool_info in tools_response.tools:
                    # 创建工具包装器
                    def create_tool_wrapper(tool_name):
                        def tool_func(tool_input):
                            # 处理工具输入
                            if isinstance(tool_input, str):
                                # 如果是字符串，尝试解析为JSON或作为URL处理
                                try:
                                    import json
                                    kwargs = json.loads(tool_input)
                                except:
                                    # 如果不是JSON，假设是URL
                                    kwargs = {"url": tool_input}
                            elif isinstance(tool_input, dict):
                                kwargs = tool_input
                            else:
                                kwargs = {"url": str(tool_input)}
                            
                            # 同步包装器，调用异步函数
                            import asyncio
                            try:
                                loop = asyncio.get_event_loop()
                                if loop.is_running():
                                    # 如果已经在事件循环中，创建新任务
                                    import concurrent.futures
                                    with concurrent.futures.ThreadPoolExecutor() as executor:
                                        future = executor.submit(asyncio.run, self._call_tool_async(tool_name, kwargs))
                                        return future.result()
                                else:
                                    return loop.run_until_complete(self._call_tool_async(tool_name, kwargs))
                            except RuntimeError:
                                return asyncio.run(self._call_tool_async(tool_name, kwargs))
                        return tool_func
                    
                    # 创建 LangChain Tool
                    langchain_tool = Tool(
                        name=tool_info.name,
                        description=tool_info.description or f"Firecrawl工具: {tool_info.name}",
                        func=create_tool_wrapper(tool_info.name)
                    )
                    
                    self.tools.append(langchain_tool)
                    
        except Exception as e:
            print(f"设置工具时出错: {e}")
            # 如果 MCP 连接失败，创建模拟工具
            await self.create_mock_tools()
            
        print(f"工具设置完成，共 {len(self.tools)} 个工具")
        
    async def _call_tool_async(self, tool_name: str, kwargs: dict):
        """异步调用MCP工具"""
        try:
            async with self.client.session("firecrawl-mcp") as session:
                result = await session.call_tool(tool_name, kwargs)
                return result.content[0].text if result.content else "无结果"
        except Exception as e:
            return f"工具调用失败: {str(e)}"
        
        
    def setup_llm(self):
        """设置语言模型"""
        print("正在设置语言模型...")
        
        api_key = os.getenv("DASHSCOPE_API_KEY")
        self.llm = Tongyi(
            model="qwen-turbo",
            api_key=api_key,
            temperature=0.7,
            max_tokens=2048
        )
        
        print("语言模型设置完成")
        
    def setup_agent(self):
        """设置智能体"""
        print("正在设置智能体...")
        
        # 创建标准的 ReAct 提示模板
        prompt_template = """你是一个专业的网页内容抓取助手，可以使用工具来获取网页内容。

你有以下工具可以使用:
{tools}

使用以下格式:

Question: 用户的问题
Thought: 你应该思考要做什么
Action: 要采取的行动，应该是 [{tool_names}] 中的一个
Action Input: 行动的输入
Observation: 行动的结果
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始问题的最终答案

开始!

Question: {input}
Thought:{agent_scratchpad}"""
        
        prompt = PromptTemplate(
            template=prompt_template,
            input_variables=["tools", "tool_names", "input", "agent_scratchpad"]
        )
        
        # 创建 React 智能体
        agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )
        
        # 创建智能体执行器
        self.agent = AgentExecutor(
            agent=agent,
            tools=self.tools,
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=5,
            return_intermediate_steps=True
        )
        
        print("智能体设置完成")
        
    async def run_query(self, query: str) -> str:
        """执行查询"""
        print(f"\n执行查询: {query}")
        print("=" * 50)
        
        try:
            # 使用智能体处理查询
            result = await self.agent.ainvoke({"input": query})
            return result["output"]
        except Exception as e:
            print(f"查询执行出错: {e}")
            # 如果智能体执行失败，直接使用工具
            if self.tools:
                tool_result = self.tools[0].func("https://book.douban.com/tag/%E5%B0%8F%E8%AF%B4")
                return f"直接工具调用结果:\n{tool_result}"
            return f"查询失败: {e}"
            
    async def initialize(self):
        """初始化智能体"""
        print("=== 初始化 Firecrawl MCP 智能体 ===")
        
        # 设置各个组件
        await self.setup_mcp_client()
        await self.setup_tools()
        self.setup_llm()
        self.setup_agent()
        
        print("\n=== 智能体初始化完成 ===")
        
async def main():
    """主函数"""
    # 创建智能体实例
    agent = FirecrawlMCPAgent()
    
    # 初始化智能体
    await agent.initialize()
    
    # 测试查询
    test_query = "帮我获取豆瓣读书当中的小说标签（https://book.douban.com/tag/%E5%B0%8F%E8%AF%B4）下的网页内容。"
    
    print(f"\n=== 测试查询 ===")
    print(f"用户输入: {test_query}")
    
    # 执行查询
    result = await agent.run_query(test_query)
    
    print(f"\n=== 查询结果 ===")
    print(result)
    
    print(f"\n=== 测试完成 ===")

if __name__ == "__main__":
    asyncio.run(main())